#include <vector>
#include <iomanip> 
#include <iostream>
#include <chrono>
#include "matrix_operations_cpu.h"
#include "simple_fnn.h"
#include "relu.h"
#include "sigmoid.h"
#include "loss_function.h"


// 生成一个指定大小和随机小数填充的矩阵
std::vector<std::vector<float>> generateRandomMatrix(int cols, int rows, float min = 0.0, float max = 1.0) {
    std::vector<std::vector<float>> matrix(rows, std::vector<float>(cols));

    // 用来生成随机数的设备
    std::random_device rd;
    // 随机数生成器
    std::mt19937 gen(rd());
    // 定义一个分布，这里使用[min, max)区间的均匀分布
    std::uniform_real_distribution<> dis(min, max);

    for (int i = 0; i < rows; ++i) {
        for (int j = 0; j < cols; ++j) {
            // 为矩阵的每个元素赋予一个随机生成的小数
            matrix[i][j] = dis(gen);
        }
    }
    
    return matrix;
}

// 函数用来打印矩阵
void printMatrix(const std::vector<std::vector<float>>& matrix) {
    for (const auto& row : matrix) {
        for (const auto& val : row) {
            std::cout << std::fixed << std::setprecision(6) << std::setw(10) << val << " ";
        }
        std::cout << std::endl;
    }
}


std::pair<std::vector<std::vector<float>>, std::vector<std::vector<float>>> generateXORData(int numSamples) {
    std::vector<std::vector<float>> inputs;
    std::vector<std::vector<float>> expectedOutputs;

    // XOR 四个基本输入和输出
    std::vector<std::vector<float>> baseInputs = {{0, 0}, {0, 1}, {1, 0}, {1, 1}};
    std::vector<std::vector<float>> baseOutputs = {{0}, {1}, {1}, {0}};

    // 循环添加数据到指定的数量
    for (int i = 0; i < numSamples; ++i) {
        int index = i % 4; // 循环索引
        inputs.push_back(baseInputs[index]);
        expectedOutputs.push_back(baseOutputs[index]);
    }

    return {inputs, expectedOutputs};
}

int main() {
    auto start = std::chrono::high_resolution_clock::now();

    auto matrixOperations = std::make_shared<MatrixOperationsCPU>(); //矩阵操作类
    auto lossFunction = std::make_shared<MeanSquaredErrorLoss>(); //损失函数类
    auto activationFunction = std::make_shared<Sigmoid>();

    // 构建神经网络的结构
    size_t inputSize = 2; // 输入层大小
    size_t hiddenLayerSize = 1024; // 隐藏层大小
    size_t numHiddenLayers = 4; //隐藏层数量
    size_t outputSize = 1; // 输出层大小
    float learningRate = 0.001;


    /**
     * 构造函数初始化一个全连接神经网络 (FNN)
     * 
     * @param inputSize 输入层的大小，即输入向量的维度。
     * @param hiddenLayerSize 隐藏层的大小，所有隐藏层都将使用这个大小。
     * @param numHiddenLayers 隐藏层的数量，定义了网络中隐藏层的总数。
     * @param outputSize 输出层的大小，即输出向量的维度。
     * @param lr 学习率，用于在训练过程中更新权重。
     * @param matOp 矩阵运算的共享指针，用于执行所有矩阵相关的操作，比如前向传播和反向传播。
     * @param lossFunction 损失函数
     */
    FNN myFNN(inputSize, hiddenLayerSize, numHiddenLayers, outputSize, learningRate, activationFunction, matrixOperations, lossFunction);

    // 构建简单的训练数据集  XOR 问题的期望输出
    // auto result = generateXORData(1000);
    // std::vector<std::vector<float>> inputs = result.first;
    // std::vector<std::vector<float>> expectedOutputs = result.second;
    std::vector<std::vector<float>> inputs = {{0, 0}, {0, 1}, {1, 0}, {1, 1}};
    std::vector<std::vector<float>> expectedOutputs = {{0}, {1}, {1}, {0}};

    // 训练过程
    int epochs = 100; // 训练轮数
    for (int epoch = 0; epoch < epochs; epoch++) {
        float totalLoss = 0.0;
        for (size_t j = 0; j < inputs.size(); j++) {
            myFNN.train({inputs[j]}, {expectedOutputs[j]});
            auto predictions = myFNN.predict({inputs[j]});
            totalLoss += myFNN.lossFunction->loss(predictions, {expectedOutputs[j]});
        }
        totalLoss /= inputs.size();
        std::cout << "Epoch " << (epoch + 1) << "/" << epochs << ", Loss: " << totalLoss << std::endl;
    }

    // 测试神经网络的表现
    for (size_t i = 0; i < inputs.size(); i++) {
        if(i > 10) break;
        auto output = myFNN.predict({inputs[i]});
        // Apply threshold
        float binaryOutput = output[0][0] >= 0.5 ? 1 : 0;
        std::cout << "Input: " << inputs[i][0] << ", " << inputs[i][1]
                << " - Expected Output: " << expectedOutputs[i][0]
                << " - Binary Output: " << binaryOutput
                << " - Predicted Output: " << output[0][0] << std::endl;
    }

    auto stop = std::chrono::high_resolution_clock::now();
    auto duration = std::chrono::duration_cast<std::chrono::milliseconds>(stop - start);
    std::cout << "Time taken by function: " << duration.count() << " milliseconds" << std::endl;
    fflush(stdout);
    return 0;
}
